package org.basis.algorithm.graph;

import org.basis.algorithm.graph.common.Edge;
import org.basis.algorithm.graph.common.Graph;
import org.basis.algorithm.graph.common.Node;

import java.util.*;

/**
 * 求图的最小生成树
 *
 * @author Mr_wenpan@163.com 2021/12/23 15:30
 */
public class Kruskal {

    public static void main(String[] args) {


    }

    public static Set<Edge> kruskalMST(Graph graph) {
        // 先构建并查集，每个点单独为一个集合
        UnionFind unionFind = new UnionFind();
        unionFind.makeSets(graph.nodes.values());

        // 将边按照权值大小排序（放入优先队列，小根堆）
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new EdgeComparator());
        for (Edge edge : graph.edges) {
            priorityQueue.add(edge);
        }
        Set<Edge> result = new HashSet<>();

        // 每次拿一条最小的边出来，检查该边的from和to节点是否在同一个集合，如果在同一个集合则说明加上这条边后会形成环，所以不能要
        // 如果没有在同一个集合，则说明该条边加入后不会形成环，可以要。然后将from和to节点所在的集合合并
        while (!priorityQueue.isEmpty()) {
            Edge edge = priorityQueue.poll();
            if (!unionFind.isSameSet(edge.from, edge.to)) {
                result.add(edge);
                unionFind.union(edge.from, edge.to);
            }
        }
        return result;
    }

    public static class EdgeComparator implements Comparator<Edge> {

        @Override
        public int compare(Edge o1, Edge o2) {
            return o1.weight - o2.weight;
        }

    }

    /**
     * 并查集
     */
    public static class UnionFind {
        private HashMap<Node, Node> fatherMap;
        private HashMap<Node, Integer> rankMap;

        public UnionFind() {
            fatherMap = new HashMap<>();
            rankMap = new HashMap<>();
        }

        private Node findFather(Node n) {
            Node father = fatherMap.get(n);
            if (father != n) {
                father = findFather(father);
            }
            fatherMap.put(n, father);
            return father;
        }

        public void makeSets(Collection<Node> nodes) {
            fatherMap.clear();
            rankMap.clear();
            for (Node node : nodes) {
                fatherMap.put(node, node);
                rankMap.put(node, 1);
            }
        }

        public boolean isSameSet(Node a, Node b) {
            return findFather(a) == findFather(b);
        }

        public void union(Node a, Node b) {
            if (a == null || b == null) {
                return;
            }
            Node aFather = findFather(a);
            Node bFather = findFather(b);
            if (aFather != bFather) {
                int aFrank = rankMap.get(aFather);
                int bFrank = rankMap.get(bFather);
                if (aFrank <= bFrank) {
                    fatherMap.put(aFather, bFather);
                    rankMap.put(bFather, aFrank + bFrank);
                } else {
                    fatherMap.put(bFather, aFather);
                    rankMap.put(aFather, aFrank + bFrank);
                }
            }
        }
    }

}
